# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
推导式发射器

处理列表推导式、集合推导式、字典推导式和生成器表达式的代码生成。
"""

from ...ir.expressions.comprehension import (
    IRListComp,
    IRSetComp,
    IRDictComp,
)
from ..core.emitter import NodeTypeEmitter, EmitResult
from ..core.context import CodeGenContext
from ...ir.expressions.basic import IRName
from ...ir.expressions.literals import IRTuple, IRList
from ..core.generator import get_global_generator


class ComprehensionEmitter(NodeTypeEmitter):
    """推导式发射器

    处理各种推导式的代码生成：
    - 列表推导式：[x*2 for x in nums]
    - 集合推导式：{x*2 for x in nums}
    - 字典推导式：{k: v for k, v in pairs}
    - 生成器表达式：(x*2 for x in nums) - 暂不支持
    """

    def __init__(self):
        super().__init__([IRListComp, IRSetComp, IRDictComp])

    def emit(self, node, context: CodeGenContext) -> EmitResult:
        """生成推导式代码"""
        if isinstance(node, IRListComp):
            return self._emit_list_comp(node, context)
        elif isinstance(node, IRSetComp):
            return self._emit_set_comp(node, context)
        elif isinstance(node, IRDictComp):
            return self._emit_dict_comp(node, context)

        return EmitResult(
            code=f"/* Unsupported comprehension type: {type(node).__name__} */",
            precedence=0,
        )

    def _emit_list_comp(self, node: IRListComp, context: CodeGenContext) -> EmitResult:
        """生成列表推导式代码

        使用 mcpy::list_comprehension 辅助函数生成简洁代码：
        [x*2 for x in nums] -> mcpy::list_comprehension<int>(nums, [&](const auto& x) { return x*2; })
        [x for x in nums if x>0] -> mcpy::list_comprehension_if<int>(nums, ..., ...)
        """

        generator = get_global_generator()
        context.require_include("mcpy/builtins.h")

        # 获取目标元素类型
        element_type = self._get_element_type_from_context(node, context)

        # 判断是否需要显式指定元素类型
        # 当容器元素是 mc::variant 或需要类型转换时，显式指定类型
        need_explicit_type = self._needs_explicit_element_type(
            node, context, element_type
        )

        # 检查是否可以使用简化的辅助函数
        if len(node.generators) == 1:
            gen = node.generators[0]
            iter_code = generator.generate_expression(gen.iter, context)

            # 获取循环变量
            if isinstance(gen.target, IRName):
                var_name = gen.target.id
            else:
                var_names = self._extract_var_names(gen.target)
                var_name = f"[{', '.join(var_names)}]"

            # 生成元素表达式
            elt_code = generator.generate_expression(node.elt, context)

            # 检查是否有条件
            if not gen.ifs:
                # 无条件：使用 list_comprehension
                if need_explicit_type:
                    code = f"mcpy::list_comprehension<{element_type}>({iter_code}, [&](const auto& {var_name}) {{ return {elt_code}; }})"
                else:
                    code = f"mcpy::list_comprehension({iter_code}, [&](const auto& {var_name}) {{ return {elt_code}; }})"
            elif len(gen.ifs) == 1:
                # 单条件：使用 list_comprehension_if
                if_code = generator.generate_expression(gen.ifs[0], context)
                if need_explicit_type:
                    code = (
                        f"mcpy::list_comprehension_if<{element_type}>({iter_code}, "
                        f"[&](const auto& {var_name}) {{ return {elt_code}; }}, "
                        f"[&](const auto& {var_name}) {{ return {if_code}; }})"
                    )
                else:
                    code = (
                        f"mcpy::list_comprehension_if({iter_code}, "
                        f"[&](const auto& {var_name}) {{ return {elt_code}; }}, "
                        f"[&](const auto& {var_name}) {{ return {if_code}; }})"
                    )
            else:
                # 多条件：使用展开代码
                code = self._emit_list_comp_expanded(node, context)
        else:
            # 嵌套生成器：使用展开代码
            code = self._emit_list_comp_expanded(node, context)

        return EmitResult(code=code, precedence=0)

    def _emit_list_comp_expanded(
        self, node: IRListComp, context: CodeGenContext
    ) -> str:
        """生成展开的列表推导式代码（用于复杂情况）"""

        generator = get_global_generator()

        # 获取元素类型
        element_type = self._get_element_type_from_context(node, context)

        lines = []

        # 使用 list_comp_nested 辅助函数包装
        # 这样可以自动进行类型转换
        lines.append(f"mcpy::list_comp_nested<{element_type}>([&]() {{")
        lines.append("    mc::variants __result;")

        # 生成嵌套的 for 循环
        indent = "    "
        for gen in node.generators:
            iter_code = generator.generate_expression(gen.iter, context)

            if isinstance(gen.target, IRName):
                var_name = gen.target.id
                lines.append(f"{indent}for (const auto& {var_name} : {iter_code}) {{")
            else:
                var_names = self._extract_var_names(gen.target)
                var_list = ", ".join(var_names)
                lines.append(f"{indent}for (const auto& [{var_list}] : {iter_code}) {{")

            indent += "    "

            # 生成 if 条件
            for if_expr in gen.ifs:
                if_code = generator.generate_expression(if_expr, context)
                lines.append(f"{indent}if ({if_code}) {{")
                indent += "    "

        # 生成元素表达式并添加到结果
        elt_code = generator.generate_expression(node.elt, context)
        lines.append(f"{indent}__result.push_back({elt_code});")

        # 关闭所有嵌套的块
        for gen in node.generators:
            for _ in gen.ifs:
                indent = indent[:-4]
                lines.append(f"{indent}}}")
            indent = indent[:-4]
            lines.append(f"{indent}}}")

        lines.append("    return __result;")
        lines.append("})")  # 关闭 lambda 和 list_comp_nested 调用

        return "\n".join(lines)

    def _get_element_type_from_context(
        self, node: IRListComp, context: CodeGenContext
    ) -> str:
        """从上下文获取列表元素类型

        简化版本：直接从节点的 cpp_type 提取元素类型。
        类型推导阶段已经利用期望类型机制正确设置了推导式的类型，
        不再需要回溯父节点。

        对于 list[int]，cpp_type 是 mc::array<int>，提取出 int。
        对于 list[Any]，cpp_type 是 mc::array<mc::variant>，提取出 mc::variant。
        """
        # 从推导式节点本身获取 C++ 类型
        cpp_type = context.get_cpp_type(node)

        if cpp_type and cpp_type.startswith("mc::array<") and cpp_type.endswith(">"):
            # 提取元素类型：mc::array<int> -> int
            element_type = cpp_type[10:-1]  # 去掉 "mc::array<" 和 ">"
            # 验证提取的类型有效（不是空字符串或占位符）
            if element_type and not element_type.startswith("{"):
                return element_type

        # 回退：从元素表达式推导类型
        elt_type = context.get_cpp_type(node.elt)
        if elt_type and elt_type != "mc::variant":
            return elt_type

        # 默认使用 mc::variant
        return "mc::variant"

    def _needs_explicit_element_type(
        self, node: IRListComp, context: CodeGenContext, element_type: str
    ) -> bool:
        """判断是否需要显式指定元素类型

        当以下情况需要显式指定：
        1. 迭代的容器元素是 mc::variant（即 list[Any]），且目标类型不是 mc::variant
        2. 目标元素类型是 mc::variant，但源容器元素不是 mc::variant
        3. 源容器元素是 mc::variant，但目标类型不是 mc::variant

        关键：必须检查**源容器**的元素类型，而不是循环变量或元素表达式的类型，
        因为类型推导可能已经将它们推导为期望类型。
        """
        if len(node.generators) == 0:
            return False

        gen = node.generators[0]

        # 获取迭代源的 inferred_type，检查其容器元素类型
        from ...type_system import TypeKind

        iter_inferred_type = gen.iter.get_extension_param("inferred_type")
        source_element_is_variant = False

        # 检查迭代源容器的元素类型是否是 ANY
        if (
            iter_inferred_type
            and hasattr(iter_inferred_type, "args")
            and iter_inferred_type.args
        ):
            # 容器的第一个类型参数是元素类型
            source_elem_type = iter_inferred_type.args[0]
            source_element_is_variant = source_elem_type.kind == TypeKind.ANY

        # 如果没有 inferred_type，检查 cpp_type
        if not source_element_is_variant:
            iter_cpp_type = context.get_cpp_type(gen.iter)
            if iter_cpp_type:
                source_element_is_variant = "mc::variant" in iter_cpp_type

        # 判断是否需要显式类型转换：
        # 1. 从 ANY 容器转换到具体类型容器：需要类型转换
        if source_element_is_variant and element_type != "mc::variant":
            return True

        # 2. 从具体类型容器转换到 ANY 容器：需要类型提升
        if not source_element_is_variant and element_type == "mc::variant":
            return True

        return False

    def _emit_set_comp(self, node: IRSetComp, context: CodeGenContext) -> EmitResult:
        """生成集合推导式代码

        使用 mcpy::set_comprehension 辅助函数生成简洁代码：
        {x*2 for x in nums} -> mcpy::set_comprehension(nums, [&](const auto& x) { return x*2; })
        {x for x in nums if x>0} -> mcpy::set_comprehension_if(nums, ..., ...)
        """

        generator = get_global_generator()
        context.require_include("mcpy/builtins.h")

        # 检查是否可以使用简化的辅助函数
        if len(node.generators) == 1:
            gen = node.generators[0]
            iter_code = generator.generate_expression(gen.iter, context)

            # 获取循环变量
            if isinstance(gen.target, IRName):
                var_name = gen.target.id
            else:
                var_names = self._extract_var_names(gen.target)
                var_name = f"[{', '.join(var_names)}]"

            # 生成元素表达式
            elt_code = generator.generate_expression(node.elt, context)

            # 检查是否有条件
            if not gen.ifs:
                # 无条件：使用 set_comprehension
                code = f"mcpy::set_comprehension({iter_code}, [&](const auto& {var_name}) {{ return {elt_code}; }})"
            elif len(gen.ifs) == 1:
                # 单条件：使用 set_comprehension_if
                if_code = generator.generate_expression(gen.ifs[0], context)
                code = (
                    f"mcpy::set_comprehension_if({iter_code}, "
                    f"[&](const auto& {var_name}) {{ return {elt_code}; }}, "
                    f"[&](const auto& {var_name}) {{ return {if_code}; }})"
                )
            else:
                # 多条件：使用展开代码
                code = self._emit_set_comp_expanded(node, context)
        else:
            # 嵌套生成器：使用展开代码
            code = self._emit_set_comp_expanded(node, context)

        return EmitResult(code=code, precedence=0)

    def _emit_set_comp_expanded(self, node: IRSetComp, context: CodeGenContext) -> str:
        """生成展开的集合推导式代码（用于复杂情况）"""

        generator = get_global_generator()

        # 获取元素类型
        element_type = self._get_element_type_from_set_context(node, context)

        lines = []

        # 使用 set_comp_nested 辅助函数包装
        lines.append(f"mcpy::set_comp_nested<{element_type}>([&]() {{")
        lines.append("    mc::variants __result;")

        # 生成嵌套的 for 循环
        indent = "    "
        for gen in node.generators:
            iter_code = generator.generate_expression(gen.iter, context)

            if isinstance(gen.target, IRName):
                var_name = gen.target.id
                lines.append(f"{indent}for (const auto& {var_name} : {iter_code}) {{")
            else:
                var_names = self._extract_var_names(gen.target)
                var_list = ", ".join(var_names)
                lines.append(f"{indent}for (const auto& [{var_list}] : {iter_code}) {{")

            indent += "    "

            # 生成 if 条件
            for if_expr in gen.ifs:
                if_code = generator.generate_expression(if_expr, context)
                lines.append(f"{indent}if ({if_code}) {{")
                indent += "    "

        # 生成元素表达式并添加到结果
        elt_code = generator.generate_expression(node.elt, context)
        lines.append(f"{indent}__result.push_back({elt_code});")

        # 关闭所有嵌套的块
        for gen in node.generators:
            for _ in gen.ifs:
                indent = indent[:-4]
                lines.append(f"{indent}}}")
            indent = indent[:-4]
            lines.append(f"{indent}}}")

        lines.append("    return __result;")
        lines.append("})")  # 关闭 lambda 和 set_comp_nested 调用

        return "\n".join(lines)

    def _get_element_type_from_set_context(
        self, node: IRSetComp, context: CodeGenContext
    ) -> str:
        """从上下文获取集合元素类型

        对于 set[int]，cpp_type 是 std::set<int>，提取出 int。
        对于 set[Any]，cpp_type 是 std::set<mc::variant>，提取出 mc::variant。
        """
        # 从推导式节点本身获取 C++ 类型
        cpp_type = context.get_cpp_type(node)

        if cpp_type and cpp_type.startswith("std::set<") and cpp_type.endswith(">"):
            # 提取元素类型：std::set<int> -> int
            element_type = cpp_type[9:-1]  # 去掉 "std::set<" 和 ">"
            # 验证提取的类型有效（不是空字符串或占位符）
            if element_type and not element_type.startswith("{"):
                return element_type

        # 回退：从元素表达式推导类型
        elt_type = context.get_cpp_type(node.elt)
        if elt_type and elt_type != "mc::variant":
            return elt_type

        # 默认使用 mc::variant
        return "mc::variant"

    def _emit_dict_comp(self, node: IRDictComp, context: CodeGenContext) -> EmitResult:
        """生成字典推导式代码

        使用 mcpy::dict_comprehension 辅助函数生成简洁代码：
        {k: v for x in items} -> mcpy::dict_comprehension(items, [&](const auto& x) { return std::make_pair(k, v); })
        {k: v for x in items if cond} -> mcpy::dict_comprehension_if(items, ..., ...)
        """

        generator = get_global_generator()
        context.require_include("mcpy/builtins.h")
        context.require_include("<utility>")  # for std::make_pair

        # 检查是否可以使用简化的辅助函数
        if len(node.generators) == 1:
            gen = node.generators[0]
            iter_code = generator.generate_expression(gen.iter, context)

            # 获取循环变量
            if isinstance(gen.target, IRName):
                var_name = gen.target.id
            else:
                var_names = self._extract_var_names(gen.target)
                var_name = f"[{', '.join(var_names)}]"

            # 生成键值表达式
            key_code = generator.generate_expression(node.key, context)
            value_code = generator.generate_expression(node.value, context)
            pair_expr = f"std::make_pair({key_code}, {value_code})"

            # 检查是否有条件
            if not gen.ifs:
                # 无条件：使用 dict_comprehension
                code = f"mcpy::dict_comprehension({iter_code}, [&](const auto& {var_name}) {{ return {pair_expr}; }})"
            elif len(gen.ifs) == 1:
                # 单条件：使用 dict_comprehension_if
                if_code = generator.generate_expression(gen.ifs[0], context)
                code = (
                    f"mcpy::dict_comprehension_if({iter_code}, "
                    f"[&](const auto& {var_name}) {{ return {pair_expr}; }}, "
                    f"[&](const auto& {var_name}) {{ return {if_code}; }})"
                )
            else:
                # 多条件：使用展开代码
                code = self._emit_dict_comp_expanded(node, context)
        else:
            # 嵌套生成器：使用展开代码
            code = self._emit_dict_comp_expanded(node, context)

        return EmitResult(code=code, precedence=0)

    def _emit_dict_comp_expanded(
        self, node: IRDictComp, context: CodeGenContext
    ) -> str:
        """生成展开的字典推导式代码（用于复杂情况）"""

        generator = get_global_generator()

        # 获取 key 和 value 类型
        key_type, value_type = self._get_dict_types_from_context(node, context)

        lines = []

        # 使用 dict_comp_nested 辅助函数包装
        lines.append(f"mcpy::dict_comp_nested<{key_type}, {value_type}>([&]() {{")
        lines.append("    mc::dict __result;")

        # 生成嵌套的 for 循环
        indent = "    "
        for gen in node.generators:
            iter_code = generator.generate_expression(gen.iter, context)

            if isinstance(gen.target, IRName):
                var_name = gen.target.id
                lines.append(f"{indent}for (const auto& {var_name} : {iter_code}) {{")
            else:
                var_names = self._extract_var_names(gen.target)
                var_list = ", ".join(var_names)
                lines.append(f"{indent}for (const auto& [{var_list}] : {iter_code}) {{")

            indent += "    "

            # 生成 if 条件
            for if_expr in gen.ifs:
                if_code = generator.generate_expression(if_expr, context)
                lines.append(f"{indent}if ({if_code}) {{")
                indent += "    "

        # 生成键值表达式并添加到结果
        key_code = generator.generate_expression(node.key, context)
        value_code = generator.generate_expression(node.value, context)
        lines.append(f"{indent}__result[{key_code}] = {value_code};")

        # 关闭所有嵌套的块
        for gen in node.generators:
            for _ in gen.ifs:
                indent = indent[:-4]
                lines.append(f"{indent}}}")
            indent = indent[:-4]
            lines.append(f"{indent}}}")

        lines.append("    return __result;")
        lines.append("})")  # 关闭 lambda 和 dict_comp_nested 调用

        return "\n".join(lines)

    def _get_dict_types_from_context(
        self, node: IRDictComp, context: CodeGenContext
    ) -> tuple[str, str]:
        """从上下文获取字典的 key/value 类型

        注意：mc::dict 不是 C++ 泛型模板类型，它总是存储 mc::variant。
        Python 的 dict[int, str] 在 C++ 中是 mc::dict（不是 mc::dict<int, std::string>）。

        这个方法从 key/value 表达式推导出它们的 C++ 类型，
        用于 dict_comp_nested 辅助函数中的类型转换逻辑。
        """
        # 直接从 key 和 value 表达式推导 C++ 类型
        key_type = context.get_cpp_type(node.key)
        value_type = context.get_cpp_type(node.value)

        # 默认使用 mc::variant 作为回退
        if not key_type or key_type == "mc::variant":
            key_type = "mc::variant"
        if not value_type or value_type == "mc::variant":
            value_type = "mc::variant"

        return (key_type, value_type)

    def _extract_var_names(self, target) -> list[str]:
        """从目标中提取变量名（支持元组解包）"""

        names = []
        if isinstance(target, IRName):
            names.append(target.id)
        elif isinstance(target, (IRTuple, IRList)):
            for element in target.elements:
                names.extend(self._extract_var_names(element))

        return names

    def get_priority(self) -> int:
        """推导式发射器的优先级"""
        return 100  # 中等优先级
